import os
import pandas as pd
import sys
import torch
from threading import Thread
from transformers import TextIteratorStreamer
from loguru import logger
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
sys.path.append(os.path.abspath('../LLM_APIs'))
from response_process import eval_response, process_llm_history


def build_local_prompt(**kwargs):
    try:
        task = kwargs['paras']['task']
        query = kwargs['paras']['query']
        texts = kwargs['paras']['texts']
        local_llm= kwargs['paras']['local_model']
        local_llm_tz = kwargs['paras']['local_tz']
    except Exception as e:
        print('Bad paras inputs... because: {}'.format(e))
        return None
    
    his_record = process_llm_history(kwargs['config'], kwargs['paras'])
    temperature = 0.01
    max_tokens = 500
    top_k = 5
    
    if task=='rewrite-paras': # when generating response using KB
        max_tokens = kwargs['paras']['words']
        topic = kwargs['paras']['topic']
        avoid_topics = kwargs['paras']['avoid_topics']
        doc_type = kwargs['paras']['type']
        style = kwargs['paras']['style']
        temperature = 0.2
        
        prompt =  f"""
                你是{doc_type}方面的撰写专家，你的任务是撰写一段主题为'''{topic}'''的专业内容。具体要求如下：
                - 撰写过程可以参考以下基础材料：
                   '''
                   {texts}
                   '''
                - 如果你认为基础材料不相关，不必参考其内容。
                - 如果材料不足，你可以自行补充，保证字数在{int(1.2*max_tokens)}字左右。
                - 禁止撰写'''{avoid_topics}'''等其他主题内容。
                - 可以采取分段论述，也可使用1)/ a)/1. 等分小点论述。
                - 不要使用MARKDOWN符号。
                - 除撰写的内容外，不要输出其他解释。
            """

    elif task=='rewrite-paras-whole':
        max_tokens = kwargs['paras']['out_limit']
        add_req = kwargs['paras']['add_req']
        
        prompt =  f"""
                你的任务是，润色并重新生成以下段落，不少于{int(1.2*max_tokens)}字，不超过{int(2*max_tokens)}字：
                '''{texts}'''
                
                生成过程需衔接和延续上文内容（如有）]：
                '''{query}'''

                生成内容要使用 1. a. 1） 等形式，分层分点论述，分点论述形式也要延续上文形式。
                禁止重复上文内容，但数字形式、公式等均与参考上文一致。
                生成内容不要使用MARKDOWN符号。
                {add_req}
            """
              
    elif task=='talk-kb':
        add_req = kwargs['paras']['add_req']
        max_tokens = kwargs['paras']['out_limit']
        prompt = f"""
                你将接收到用户的问题：'''{query}'''
                你将接收到一些基础知识和资料：'''{texts}'''
                你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                请严格参考基础知识资料以及历史记录（如有），回答用户问题。{add_req}
                如果用户在找图片，且基础资料中存在该图片，直接返回图片的名字。
            """
            
    elif task=='summary':
        try:
            limit = kwargs['paras']['out_limit']
            add_req = kwargs['paras']['add_req']
        except:
            limit = int(len(texts)*0.2)
            add_req = ''
        prompt = f"""
                你将接收一段文字材料：{texts}
                你需要总结以上文字材料主要内容，字数不超过{limit}字。{add_req}
                你的回答只包含提炼后的文本，不要返回其他多余信息。
            """

    elif task=='summary-plain':
        max_tokens = kwargs['paras']['out_limit']
        add_req = kwargs['paras']['add_req']

        prompt = f"""
                对下面这段文本总结提炼为一段精简的摘要，保留其关键含义，不要超过{max_tokens}字。{add_req}
                '''{texts}'''
            """
        
    elif task=='if-history':
        max_tokens = 10
        example_true = '{"trace_history":"是"}'
        example_false = '{"trace_history":"否"}'
        examples_full = '''
                        上一轮提问是：360安全大模型的技术方案包括哪些内容？
                        当前提问是：其中第二项技术方案的具体实现方式是是什么？
                        你应该判断当前提问是上一轮提问的追问，因此回答 {"trace_history":"是"}
                        
                        上一轮提问是：基坑旁边是否可以存放施工材料？
                        当前提问是：室内管道的最低标高要求是多少？
                        你应该判断当前提问不是上一轮提问的追问，因此回答 {"trace_history":"否"}
        '''
        histories = kwargs['paras']['histories']
        if len(histories)>0:
            last_query = histories[-1][0]
        else:
            last_query = ''
            
        prompt = f"""
                你会收到当前用户的提问'''{query}'''
                你还会收到上一轮用户的提问'''{last_query}'''
                请判断，当前提问'''{query}'''是否为在上一轮提问'''{last_query}'''基础上的追加问题
                你的回答为JSON字典格式，键为"trace_history"，值只能为"是"或者"否"，例如{example_true}或者{example_false}
                不要输出其他内容，回答不要超过5个字。下面给出了一些例子：
            """
        prompt = prompt + examples_full

    elif task=='gen-root-titles':
        content = kwargs['paras']['content']
        style = kwargs['paras']['style']
        title = kwargs['paras']['title']
        doc_type = kwargs['paras']['doc_type']
        num_titles = kwargs['paras']['num_titles']
        user_thoughts = kwargs['paras']['thoughts']
        max_tokens = 500
                                
        prompt = f"""你是{doc_type}文件撰写方面的专家，需要你生成文件提纲，文件标题是{title}。具体要求如下：
                - 生成的投标文件提纲要参考以下招标要求\n：{content}\n
                - 仅生成一级提纲（即最高层级的章节标题），标题数量为{num_titles}个。
                - 生成的提纲层级遵循Markdown格式，即所有标题都使用"-"前缀分隔。
                - 生成的风格符合{style}，除生成提纲之外不要返回任何其他解释。
                - 生成提纲可以参考以下样例：
                {texts}
                - 需要注意用户思路（如有）：{user_thoughts}
            """
        
    elif task=='gen-titles-oneoff':
        topic = kwargs['paras']['topic']
        doc_type = kwargs['paras']['doc_type']
        num_subs = kwargs['paras']['num_subs']
        max_depth = kwargs['paras']['max_depth']
        max_tokens = kwargs['paras']['max_tokens']
        num_titles = kwargs['paras']['num_titles']
                                
        prompt = f"""你是{doc_type}文件撰写方面的专家，需要你以{topic}为文件主题，生成完整的文件提纲，具体要求如下：
                - 生成的提纲要具备清晰完整的父-子树形层级结构，层级数为{max_depth} ，每层包含的标题数量不超过{num_subs}个。
                - 生成的提纲层级遵循Markdown格式，即所有标题都使用"-"前缀，并使用"-"和空格区分标题层级。
                - 全部标题数量不少于{num_titles}个，不要生成意思重复的标题，除生成提纲之外不要返回任何其他解释。
                - 提纲内容需参考和包括以下内容\n:
                {texts}
            """

    elif task=='filling-tb': # when generating response using KB
        temperature = 0.01
        prompt = f"""
                你是一名表格填充人员，你将接收基础资料\n：'''{texts}'''
                请根据已有资料，找到对应查询关键词'''{query}'''的答案。注意以下事项：
                - 如果基础资料中没有包括答案，你的答案是"null"。
                - 注意选择题和打☑的项。
                - 必须返回JSON字典格式，键为"answer"，值为你找到的答案。
                - 答案要直接简单，不要添加任何解释，不要输出任何其他信息。
            """
 
    elif task=='extract-contents':
        doc_type = kwargs['paras']['doc_type']
        topic = kwargs['paras']['topic']
        avoid_topics = kwargs['paras']['avoid_topics']
        
        prompt = f"""你是{doc_type}文件撰写方面的专家，需要分析以下内容，仅提取与{topic}有关的内容，具体要求如下：
                '''
                {texts}
                '''
                - 仅提取与{topic}相关的内容，不要提取{avoid_topics}方面的内容。
                - 不要输出额外解释。
            """
        
    elif task=='extract-topics':
        doc_type = kwargs['paras']['doc_type']
        num_topics = kwargs['paras']['num_topics']

        prompt = f"""你是{doc_type}文件撰写方面的专家，需要分析以下内容，从总提取数个主题，具体要求如下：
                '''
                {texts}
                '''
                - 你自主判断提取的主题数量，但不能超过{num_topics}个。
                - 保证主题涉及的内容尽量区分，不要重叠。
                - 不要输出额外解释。
                - 输出采用Markdown格式，使用"-"和空格区分主题。
            """
    
    elif task=='summary-keywords': # can be merged with task extract-topics
        max_tokens = kwargs['paras']['out_limit']
        kw_num = kwargs['paras']['kw_number']
        add_req = kwargs['paras']['add_req']
        
        prompt = f"""
                提取下面这段文字的主题关键词，数量不超过{kw_num}个，关键词之间用分号";"隔开。{add_req}
                '''{texts}'''
            """
                  
    elif task=='plain':
        res_text = texts

    limit_len = 20000
    prompt = prompt[:limit_len] + '...' if len(prompt) > limit_len else prompt
    return prompt, local_llm, local_llm_tz, temperature, top_k, max_tokens


def _gc():
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


SYSTEM_ROLE = '你是一位得力的助手，拥有渊博的行业知识。'
def call_local_llm(**kwargs):
    prompt, local_llm, local_llm_tz, temperature, top_k, max_tokens = build_local_prompt(**kwargs)
    logger.debug('prompt_len:{}', len(prompt))
    _gc()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    messages = [{'role': 'system', 'content': SYSTEM_ROLE}, {"role": "user", "content": prompt}]
    
    try:
        inputs = local_llm_tz.apply_chat_template(messages,
                                                add_generation_prompt=True,
                                                tokenize=False,
                                                return_tensors="pt",
                                                )
        
        model_inputs = local_llm_tz([inputs], return_tensors="pt").to(device)
        generated_ids = local_llm.generate(
            model_inputs.input_ids,
            max_new_tokens=max_tokens*3,
            temperature=temperature,
            top_k=top_k,
            do_sample=True,
            pad_token_id=local_llm_tz.eos_token_id 
        )

        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]  
        res_text = local_llm_tz.batch_decode(generated_ids, skip_special_tokens=True)[0]
        res_text = eval_response(res_text)
        
    except Exception as e:
        print('\t', e)
        res_text = None
    
    logger.debug('res_text:{}', res_text)
    return res_text


def _chat_stream(model, tokenizer, query, history, temperature, top_k, max_tokens):
    conversation = [{'role': 'system', 'content': SYSTEM_ROLE}, ]
    for query_h, response_h in history:
        conversation.append({"role": "user", "content": query_h})
        conversation.append({"role": "assistant", "content": response_h})
    conversation.append({"role": "user", "content": query})
    input_text = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(
        tokenizer=tokenizer, skip_prompt=True, timeout=60.0, skip_special_tokens=True
    )
    generation_kwargs = {
        **inputs,
        "streamer": streamer,
        "temperature": temperature,
        "top_k": top_k,
        "max_new_tokens": max_tokens*2
    }
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for new_text in streamer:
        yield new_text


def call_local_stream(**kwargs):
    prompt, local_llm, local_llm_tz, temperature, top_k, max_tokens = build_local_prompt(**kwargs)
    logger.debug('prompt_len:{}', len(prompt))
    history = []
    _gc()
    res = ""
    for res_text in _chat_stream(local_llm, local_llm_tz, prompt, history, temperature, top_k, max_tokens):
        if res_text:
            res += res_text
            res_text = res_text.replace('\n', '</br>')
            yield f'data: {res_text}\n\n'
    logger.debug('res:{}', res)


